﻿using Fst.SemiRing;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Text;
using System.Runtime.Serialization.Formatters.Binary;

namespace Fst
{
    /// <summary>
    /// A mutable finite state transducer implementation.
    /// 
    /// Holds an ArrayList of {@link edu.cmu.sphinx.fst.State} objects allowing
    /// additions/deletions.
    /// 
    /// @author John Salatas <jsalatas@users.sourceforge.net>
    /// </summary>
    public class Fst
    {
        // fst states
        private List<State> states = null;

        // initial state
        protected State start;

        // input symbols map
        protected String[] isyms;

        // output symbols map
        protected String[] osyms;

        // semiring
        protected Semiring _semiring;

        /**
        /// Default Constructor
         */
        public Fst() {
            states = new List<State>();
        }

        /**
        /// Constructor specifying the initial capacity of the states ArrayList (this
        /// is an optimization used in various operations)
        /// 
        /// @param numStates
        ///            the initial capacity
         */
        public Fst(int numStates) 
        {
            if (numStates > 0) {
                states = new List<State>(numStates);
            }
        }

        /**
        /// Constructor specifying the fst's semiring
        /// 
        /// @param s
        ///            the fst's semiring
         */
        public Fst(Semiring s) 
            :this()
        {
            this._semiring = s;
        }

        /**
        /// Get the initial states
         */
        public State getStart() 
        {
            return start;
        }

        /**
        /// Get the semiring
         */
        public Semiring getSemiring() 
        {
            return _semiring;
        }

        /**
        /// Set the Semiring
        /// 
        /// @param semiring
        ///            the semiring to set
         */
        public void setSemiring(Semiring semiring) 
        {
            this._semiring = semiring;
        }

        /**
        /// Set the initial state
        /// 
        /// @param start
        ///            the initial state
         */
        public void setStart(State start) 
        {
            this.start = start;
        }

        /**
        /// Get the number of states in the fst
         */
        public int getNumStates() 
        {
            return this.states.Count;
        }

        public State getState(int index) 
        {
            return states[index];
        }

        /**
        /// Adds a state to the fst
        /// 
        /// @param state
        ///            the state to be added
         */
        public void addState(State state) 
        {
            this.states.Add(state);
            state.id = states.Count - 1;
        }

        /**
        /// Get the input symbols' array
         */
        public String[] getIsyms() 
        {
            return isyms;
        }

        /**
        /// Set the input symbols
        /// 
        /// @param isyms
        ///            the isyms to set
         */
        public void setIsyms(String[] isyms) 
        {
            this.isyms = isyms;
        }

        /**
        /// Get the output symbols' array
         */
        public String[] getOsyms() 
        {
            return osyms;
        }

        /**
        /// Set the output symbols
        /// 
        /// @param osyms
        ///            the osyms to set
         */
        public void setOsyms(String[] osyms) 
        {
            this.osyms = osyms;
        }

        /**
        /// Serializes a symbol map to an ObjectOutputStream
        /// 
        /// @param out
        ///            the ObjectOutputStream. It should be already be initialized by
        ///            the caller.
        /// @param map
        ///            the symbol map to serialize
        /// @throws IOException
         */
        private void writeStringMap(StreamWriter OutStream,String[] map)
        {
            OutStream.Write(map.Length);
            for (int i = 0; i < map.Length; i++) 
            {
                OutStream.Write(map[i]);
            }
        }

        /**
        /// Serializes the current Fst instance to an ObjectOutputStream
        /// 
        /// @param out
        ///            the ObjectOutputStream. It should be already be initialized by
        ///            the caller.
        /// @throws IOException
         */
        private void writeFst(StreamWriter OutStream) 
        {
            writeStringMap(OutStream,isyms);
            writeStringMap(OutStream,osyms);
            OutStream.Write(states.IndexOf(start));

            OutStream.Write(_semiring);
            OutStream.Write(states.Count);

            Dictionary<State, int> stateMap = new Dictionary<State, int>(states.Count);
            for (int i = 0; i < states.Count; i++) 
            {
                State s = states[i];
                OutStream.Write(s.getNumArcs());
                OutStream.Write(s.getFinalWeight());
                OutStream.Write(s.getId());
                if(stateMap.ContainsKey(s))
                    stateMap[s]=i;
                else
                    stateMap.Add(s, i);
            }

            int numStates = states.Count;
            for (int i = 0; i < numStates; i++) 
            {
                State s = states[i];
                int numArcs = s.getNumArcs();
                for (int j = 0; j < numArcs; j++) 
                {
                    Arc a = s.getArc(j);
                    OutStream.Write(a.getIlabel());
                    OutStream.Write(a.getOlabel());
                    OutStream.Write(a.getWeight());
                    OutStream.Write(stateMap[a.getNextState()]);
                }
            }
        }

        /**
        /// Saves binary model to disk
        /// 
        /// @param filename
        ///            the binary model filename
        /// @throws IOException
         */
        public void saveModel(String filename)
        {
            StreamWriter oos = new StreamWriter(filename);
            writeFst(oos);
            oos.Flush();
            oos.Close();
        }

        /**
        /// Deserializes a symbol map from an ObjectInputStream
        /// 
        /// @param in
        ///            the ObjectInputStream. It should be already be initialized by
        ///            the caller.
        /// @return the deserialized symbol map
        /// @throws IOException
        /// @throws ClassNotFoundException
         */
        protected static String[] readStringMap(FileStream inStream)
        {
            BinaryFormatter serializer = new BinaryFormatter();
            
            int mapSize = (int)serializer.Deserialize(inStream);
            String[] map = new String[mapSize];
            for (int i = 0; i < mapSize; i++) {
                String sym = (String) serializer.Deserialize(inStream);
                map[i] = sym;
            }

            return map;
        }

        /**
        /// Deserializes an Fst from an ObjectInputStream
        /// 
        /// @param in
        ///            the ObjectInputStream. It should be already be initialized by
        ///            the caller.
        /// @return
        /// @throws IOException
        /// @throws ClassNotFoundException
         */
        private static Fst readFst(FileStream inStream) 
        {
            String[] _is = readStringMap(inStream);
            String[] os = readStringMap(inStream);
            BinaryFormatter serializer = new BinaryFormatter();

            int startid = (int)serializer.Deserialize(inStream);
            Semiring semiring = (Semiring) serializer.Deserialize(inStream);
            int numStates = (int)serializer.Deserialize(inStream);
            Fst res = new Fst(numStates);
            res.isyms = _is;
            res.osyms = os;
            res._semiring = semiring;
            for (int i = 0; i < numStates; i++) 
            {
                int numArcs = (int)serializer.Deserialize(inStream);
                State s = new State(numArcs + 1);
                float f = (float)serializer.Deserialize(inStream);
                if (f == res._semiring.zero()) {
                    f = res._semiring.zero();
                } 
                else if (f == res._semiring.one()) 
                {
                    f = res._semiring.one();
                }
                s.setFinalWeight(f);
                s.id = (int)serializer.Deserialize(inStream);
                res.states.Add(s);
            }
            res.setStart(res.states[startid]);

            numStates = res.getNumStates();
            for (int i = 0; i < numStates; i++) 
            {
                State s1 = res.getState(i);
                for (int j = 0; j < s1.initialNumArcs - 1; j++) {
                    Arc a = new Arc();
                    a.setIlabel((int)serializer.Deserialize(inStream));
                    a.setOlabel((int)serializer.Deserialize(inStream));
                    a.setWeight((float)serializer.Deserialize(inStream));
                    a.setNextState(res.states[(int)serializer.Deserialize(inStream)]);
                    s1.addArc(a);
                }
            }

            return res;
        }

        /**
        /// Deserializes an Fst from disk
        /// 
        /// @param filename
        ///            the binary model filename
        /// @throws IOException
        /// @throws ClassNotFoundException
         */
        public static Fst loadModel(String filename) 
        {
            DateTime starttime = DateTime.Now;
            Fst obj;

            FileStream fis = new FileStream(filename, FileMode.Open);
            obj = readFst(fis);
            fis.Close();

            Trace.WriteLine("Load Time: "
                    + (DateTime.Now - starttime).ToString() );
            return obj;
        }

        /*
        /// (non-Javadoc)
        /// 
        /// @see java.lang.Object#equals(java.lang.Object)
         */
        override
        public Boolean Equals(Object obj) 
        {
            if (this == obj)
                return true;
            if (obj == null)
                return false;
            if (GetType() != obj.GetType())
                return false;
            Fst other = (Fst) obj;
            if (!isyms.Equals(other.isyms))
                return false;
            if (!osyms.Equals(other.osyms))
                return false;
            if (start == null) {
                if (other.start != null)
                    return false;
            } 
            else if (!start.Equals(other.start))
                return false;
            if (states == null) 
            {
                if (other.states != null)
                    return false;
            } 
            else if (!states.Equals(other.states))
                return false;
            if (_semiring == null) {
                if (other._semiring != null)
                    return false;
            } 
            else if (!_semiring.Equals(other._semiring))
                return false;
            return true;
        }

        /*
        /// (non-Javadoc)
        /// 
        /// @see java.lang.Object#toString()
         */
        override
        public String ToString() 
        {
            StringBuilder sb = new StringBuilder();
            sb.Append("Fst(start=" + start + ", isyms=" + isyms + ", osyms="
                    + osyms + ", semiring=" + _semiring + ")\n");
            int numStates = states.Count;
            for (int i = 0; i < numStates; i++) {
                State s = states[i];
                sb.Append("  " + s + "\n");
                int numArcs = s.getNumArcs();
                for (int j = 0; j < numArcs; j++) {
                    Arc a = s.getArc(j);
                    sb.Append("    " + a + "\n");
                }
            }

            return sb.ToString();
        }

        /**
        /// Deletes a state
        /// 
        /// @param state
        ///            the state to delete
         */
        public void deleteState(State state) 
        {

            if (state == start) {
                Trace.WriteLine("Cannot delete start state.");
                return;
            }

            states.Remove(state);

            foreach (State s1 in states) 
            {
                List<Arc> newArcs = new List<Arc>();
                for (int j = 0; j < s1.getNumArcs(); j++) 
                {
                    Arc a = s1.getArc(j);
                    if (!a.getNextState().Equals(state)) 
                    {
                        newArcs.Add(a);
                    }
                }
                s1.setArcs(newArcs);
            }
        }

        /**
        /// Remaps the states' ids.
        /// 
        /// States' ids are renumbered starting from 0 up to @see
        /// {@link edu.cmu.sphinx.fst.Fst#getNumStates()}
         */
        public void remapStateIds() 
        {
            int numStates = states.Count;
            for (int i = 0; i < numStates; i++) {
                states[i].id = i;
            }

        }

        public void deleteStates(HashSet<State> toDelete) 
        {

            if (toDelete.Contains(start)) 
            {
                Trace.WriteLine("Cannot delete start state.");
                return;
            }

            List<State> newStates = new List<State>();

            foreach (State s1 in states) 
            {
                if (!toDelete.Contains(s1)) 
                {
                    newStates.Add(s1);
                    List<Arc> newArcs = new List<Arc>();
                    for (int j = 0; j < s1.getNumArcs(); j++) 
                    {
                        Arc a = s1.getArc(j);
                        if (!toDelete.Contains(a.getNextState())) 
                        {
                            newArcs.Add(a);
                        }
                    }
                    s1.setArcs(newArcs);
                }
            }
            states = newStates;

            remapStateIds();
        }

    }
}
